# -*- coding: utf-8 -*-
# @Time    : 2022/12/10 18:42
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : test.py
# @Software: PyCharm
import sys

import matplotlib

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
import numpy
import pandas
import scipy.interpolate

sys.path.append(r'F:\changeworld\caes')
import pygpt.parse_gdf as pg
from shapely.geometry import LineString
from shapely.ops import snap, split, nearest_points

EPS = 1e-4


class Traj:
    class Colname:
        all_columns = index, t, x, y, z, rxy, G, Bx, By, Bz, m, q, nmacro, rmacro, ID, fEx, fEy, fEz, fBx, fBy, fBz = 'index, t, x, y, z, rxy, G, Bx, By, Bz, m, q, nmacro, rmacro, ID, fEx, fEy, fEz, fBx, fBy, fBz'.split(
            ', ')

    def __init__(self, gdf, pt_index):
        idx = list(range(len(gdf)))
        self.data = pandas.DataFrame(columns=Traj.Colname.all_columns, index=idx)
        # self.data.reindex(list(range(len(gdf))))
        self.data[self.Colname.index] = self.data.index
        self.data[self.Colname.t] = self.data[self.Colname.index].apply(lambda i: gdf[i]['p']['time'])
        for key in self.Colname.all_columns[2:]:
            self.data[key] = self.data[self.Colname.index].apply(lambda i: gdf[i]['d'][key][pt_index])
        pass


def clipLine(ls, p1, p2):
    p1 = nearest_points(ls, p1)[0]
    geom = snap(ls, p1, 1)
    parts = split(geom, p1)
    if len(parts) == 1:
        fp = parts[0]
    else:
        fp = parts[0] if parts[0].distance(p2) < parts[1].distance(p2) else parts[1]

    p2 = nearest_points(fp, p2)[0]
    geom = snap(fp, p2, 0.1)
    parts = split(geom, p2)

    if len(parts) == 1:
        res = parts[0]
    else:
        res = parts[0] if parts[0].distance(p1) < parts[1].distance(p1) else parts[1]
    return res


BIGNUM = 1e10
ideal_gdf = pg.gdftomemory(r"F:\GPTfiles\LGB\integration\experi_ideaLGB.gdf")[:-1]
# ideal_gdf = pg.gdftomemory(r"F:\GPTfiles\LGB\integration\LGB2-res\experi_ideaLGB.gdf")[:-1]

index = range(1)
gdfs = [pg.gdftomemory(r"F:\GPTfiles\LGB\integration\experi.gdf")[:-1] for i in index]

# B_data = pandas.read_csv(r"F:\GPTfiles\LGB\integration\B_griddata.txt", )
# LGB1
ds = [0.211975929001,
      0.211999122179,
      0.211991579818,
      0.211974703576,
      0.211954069270
      ]

# LGB2
# ds = [0.22986918670005300000,
#       0.22994126382221300000,
#       0.22999018711081700000,
#       0.22997460785188300000,
#       0.22987960231039200000
#       ]


# LGB3
# ds =1e-3* numpy.array((
# 179.9673857,
# 199.9949865,
# 179.9673857
# ))
# ds = [90e-3]  # 磁铁长度SYSU

z_integration_low_limit = -400e-3  # -0.15  # 积分下限
z_integration_high_limit = 1300e-3  # 0.15  # 积分上限

z_magnet_start = 0  # -279.964879e-3  # -52.53  # -45e-3  # 磁铁端部起始位置
z_magnet_end = z_magnet_start + sum(ds)
index_of_center_pt = 13


# %%
def piecewise_pts(gdf, pt_id=0):
    zstarts = [sum(ds[0:i]) for i in range(len(ds))] + [sum(ds)]
    zstarts.append(BIGNUM)
    step = 0
    pts = []
    pts_piecewise = [list() for i in range(len(zstarts) - 1)]

    for i in range(len(gdf)):
        x = gdf[i]['d']['x'][pt_id]
        z = gdf[i]['d']['z'][pt_id]

        if z > zstarts[step + 1]:
            step += 1
            if step > len(pts_piecewise) - 1:
                break
        pts_piecewise[step].append([x, z])
        pts.append([x, z])
    return pts, pts_piecewise, zstarts


def get_pts(gdf, pt_id=0):
    zstarts = [sum(ds[0:i]) for i in range(len(ds))] + [sum(ds)]
    zstarts.append(BIGNUM)

    pts = []

    for i in range(len(gdf)):
        z = gdf[i]['d']['z'][pt_id]
        x = gdf[i]['d']['x'][pt_id]

        pts.append([x, z])
    return pts


def gdf_to_DataFrame(gdf, pt_id=0) -> pandas.DataFrame:
    values = []
    for data in gdf:
        # df = pandas.DataFrame.from_dict(data['d'])
        # values.append(df.loc[pt_id])

        values.append([data['d'][key][pt_id] for key in data['d']])
    return pandas.DataFrame(values, columns=data['d'].keys())


for pt_id in [0, gdfs[0][0]['d']['x'].shape[0] - 1]:  # range(len(gdfs[0][0]['d']['x'])):
    _df = gdf_to_DataFrame(gdfs[0], pt_id)
    _df.to_csv(r'F:\GPTfiles\LGB\integration\pys\gdf_details/%d.csv' % pt_id, index=False)


def get_delta_xs(zs1: numpy.ndarray, xs1: numpy.ndarray, zs2, xs2):
    f = scipy.interpolate.interp1d(zs1, xs1, fill_value='extrapolate')
    xs1_at_zs2 = f(zs2)
    return xs1_at_zs2 - xs2


def get_int_B(gdf, pt_id=0, z_start=0, z_end=z_magnet_end):
    Bys = [each['d']['fBy'][pt_id] for each in gdf]
    mid_Bys = [numpy.mean(Bys[i:i + 2]) for i in range(len(Bys) - 1)]

    pts = get_pts(gdf, pt_id)
    pts = numpy.array(pts).reshape(-1, 2)
    filter = (pts[:, 1] < z_integration_high_limit) & (pts[:, 1] > z_start) & (pts[:, 1] < z_end)
    pts = pts[filter]

    delta2 = (numpy.diff(pts, axis=0)) ** 2
    int_B = ((delta2[:, 0] + delta2[:, 1]) ** .5 * numpy.array(mid_Bys)[filter[:-1]][:-1]).sum()

    return int_B, pts


def int_B_uniformity(gdf, z_start, z_end, z_magnet_start=z_magnet_start):
    """
    计算积分场均匀性
    :param gdf:
    :return:
    """
    intBs = []
    trajs = []
    xs_at_ref_in = []
    xs_at_ref_out = []
    xs_prim_at_ref_in = []
    # z_out = 1.7
    for i in range(len(gdf[0]['d']['x'])):
        intB, pts = get_int_B(gdf, i, z_start, z_end)
        intBs.append(intB)
        trajs.append(pts)
        xs_at_ref_in.append(
            list(LineString(pts).intersection(LineString(((-1, z_magnet_start), (1, z_magnet_start)))).coords)[0][0])
        index_z_mag_start = numpy.where(numpy.abs(pts[:, 1] - z_magnet_start) < EPS)[0][0]

        xs_prim_at_ref_in.append(
            numpy.arctan2(numpy.diff(pts[:, 0])[index_z_mag_start], numpy.diff(pts[:, 1])[index_z_mag_start])
        )

        xs_at_ref_out.append(
            list(LineString(pts).intersection(LineString(((-1, z_magnet_end), (1, z_magnet_end)))).coords)[0][0])
    print([numpy.rad2deg(each) for each in xs_prim_at_ref_in])
    return intBs, trajs, xs_at_ref_in, xs_at_ref_out, xs_prim_at_ref_in


# %%
ideal_pts, ideal_pts_piecewise, ideal_zstarts = piecewise_pts(ideal_gdf)
ideal_pts = numpy.array(ideal_pts).reshape(-1, 2)
ideal_pts = ideal_pts[(ideal_pts[:, 1] >= z_integration_low_limit) & (ideal_pts[:, 1] <= z_integration_high_limit)]
df_ideal = pandas.DataFrame(data=ideal_pts, columns=['x', 'z'])
df_ideal.to_csv("理想磁场下电子轨迹.csv", index=False)

for ii in index:
    gdf = gdfs[ii - index[0]]
    mytraj = Traj(gdf, index_of_center_pt, )
    pts, pts_piecewise, zstarts = piecewise_pts(gdf, pt_id=index_of_center_pt)
    pts = numpy.array(pts).reshape(-1, 2)

    delta_xs = (get_delta_xs(ideal_pts[:, 1], ideal_pts[:, 0], pts[:, 1], pts[:, 0], ))
    # plt.figure()
    filter = pts[:, 1] <= z_integration_high_limit
    # plt.plot(pts[:, 1][filter], delta_xs[filter])

    # %%

    intBs, trajs, xs_at_ref_in, xs_at_ref_out, xs_prim_at_ref_in = int_B_uniformity(gdf, z_integration_low_limit,
                                                                                    z_integration_high_limit)
    plt.figure()
    for traj_ in trajs:
        plt.plot(*traj_.T[::-1])
    plt.axvline(z_magnet_start)
    plt.axvline(z_magnet_end)
    plt.gca().set_aspect('equal')
    # %%
    # plt.figure()
    # plt.plot(xs_at_ref_in, intBs / intBs[10] - 1)
    # plt.show()

    # %%
    df = pandas.DataFrame(data=numpy.vstack((xs_at_ref_in, xs_at_ref_out, intBs)).T,
                          columns=["x@z=%dmm" % (z_magnet_start * 1e3), "x@z=%dmm" % (z_magnet_end * 1e3), "int B"])
    df.to_csv("积分场均匀性-%s.csv" % ii, index=False)
    # %%

    traj = LineString(pts)
    plt.figure()
    plt.plot(*numpy.array(traj.coords).T[::-1],  # label="Opera",
             label="Measured")
    plt.plot(*ideal_pts.T[::-1], label="Ideal")
    plt.title("%d" % ii)
    plt.legend()

    bounds = [LineString([[-1., zstart], [1., zstart]]) for zstart in zstarts][1:-1]
    bound_pts = []
    for bound in bounds:
        bound_pts.append(traj.intersection(bound))

    for i in range(len(pts_piecewise) - 1):
        bound_pt = list(bound_pts[i].coords)[0]
        pts_piecewise[i].append(bound_pt)
        pts_piecewise[i + 1].insert(0, bound_pt)
    lines = [LineString(pts_) for pts_ in pts_piecewise]
    step_lens = [line.length for line in lines]
    # intBs = numpy.array(Bs) * numpy.array(step_lens)
    # print("积分场大小：%.15f T m\n各台阶积分场：%s T m\nLGB内轨迹全长：%.15f m\n各台阶轨迹长度：%s m" % (
    #     intBs.sum(), intBs[:-1], sum(step_lens[:-1]), step_lens[:-1]))
    # base_color = int('1f77b4', 16)
    # for i in range(len(zstarts) - 2):
    #     plt.axvspan(zstarts[i], zstarts[i + 1], alpha=0.3, facecolor=[(.0 + 7 * i * 2.1) % 1.0,
    #                                                                   (.0 + 2 * i * 1.3) % 1.0,
    #                                                                   (.0 + 3 * i * 1.1) % 1.0])

    # plt.axis('equal')
    # plt.show()

    # %%
    df_opera = pandas.DataFrame(data=numpy.hstack([pts[filter], delta_xs[filter].reshape(-1, 1)]),
                                columns=['x', 'z', 'delta_xs'])
    csv_name = "opera磁场下电子轨迹-%d.csv" % ii
    df_opera.to_csv(csv_name, index=False)
    print("%s done!" % csv_name)
plt.show()
